﻿using Microsoft.AspNetCore.Http;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net.WebSockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using WebSocketDemo.Models;

namespace WebSocketDemo
{
    public class ChatWebsocketMiddleware
    {
        private static ConcurrentDictionary<string, WebSocket> _sockets = new ConcurrentDictionary<string, WebSocket>();

        private readonly RequestDelegate _next;

        public ChatWebsocketMiddleware(RequestDelegate next)
        {
            _next = next;
        }

        public async Task Invoke(HttpContext context)
        {
            if (!context.WebSockets.IsWebSocketRequest)
            {
                await _next.Invoke(context);
                return;
            }

            Clear();

            CancellationToken ct = context.RequestAborted;
            var currentSocket = await context.WebSockets.AcceptWebSocketAsync();
            var socketId = context.Request.Query["sid"].ToString();
            if (!_sockets.ContainsKey(socketId))
            {
                _sockets.TryAdd(socketId, currentSocket);
            }

            while (true)
            {
                if (ct.IsCancellationRequested)
                {
                    break;
                }

                var response = await ReceiveStringAsync(currentSocket, ct);
                var msg = response.FromJson<Message>();
                if (response.IsNullOrEmpty())
                {
                    if (currentSocket.State != WebSocketState.Open)
                    {
                        break;
                    }
                    continue;
                }

                foreach (var socket in _sockets)
                {
                    if (socket.Value.State != WebSocketState.Open)
                    {
                        continue;
                    }
                    if (socket.Key == msg.ToId || socket.Key == socketId)
                    {
                        await SendStringAsnyc(socket.Value, msg.ToJson(), ct);
                    }
                }
            }

            await currentSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Closing", ct);
            currentSocket.Dispose();

        }

        private void Clear()
        {
            foreach (var key in _sockets.Keys)
            {
                var sock = _sockets[key];
                if(sock.State != WebSocketState.Open)
                {
                    _sockets.TryRemove(key, out _);
                }
            }

            //for (int i = _sockets.Keys.Count; i > 0; i--)
            //{
            //    string key = _sockets.Keys.ToList()[i];
            //    if (_sockets[key].State != WebSocketState.Open)
            //    {
            //        _sockets.TryRemove(key, out _);
            //    }
            //}
        }

        private Task SendStringAsnyc(WebSocket socket, string data, CancellationToken ct)
        {
            var buffer = Encoding.UTF8.GetBytes(data);
            var segment = new ArraySegment<byte>(buffer);
            return socket.SendAsync(segment, WebSocketMessageType.Text, true, ct);
        }

        private async Task<string> ReceiveStringAsync(WebSocket socket, CancellationToken ct)
        {
            var buffer = new ArraySegment<byte>(new byte[8192]);
            using (var msg = new MemoryStream())
            {
                WebSocketReceiveResult result;
                do
                {
                    ct.ThrowIfCancellationRequested();

                    result = await socket.ReceiveAsync(buffer, ct);
                    msg.Write(buffer.Array, buffer.Offset, result.Count);
                } while (!result.EndOfMessage);

                msg.Seek(0, SeekOrigin.Begin);
                if (result.MessageType != WebSocketMessageType.Text)
                {
                    return null;
                }

                using (var reader = new StreamReader(msg, Encoding.UTF8))
                {
                    return await reader.ReadToEndAsync();
                }
            }
        }
    }
}
